package com.example.demomcp.config;

import com.fasterxml.jackson.databind.ObjectMapper;
import io.modelcontextprotocol.server.transport.WebMvcSseServerTransport;
import io.modelcontextprotocol.spec.ServerMcpTransport;

import org.springframework.ai.autoconfigure.mcp.server.McpServerProperties;
import org.springframework.boot.autoconfigure.AutoConfiguration;
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.http.HttpStatus;
import org.springframework.web.servlet.function.*;

import java.util.Optional;

@Configuration
public class MyMcpWebConfig {
    @Bean
    @ConditionalOnMissingBean
    public WebMvcSseServerTransport webMvcSseServerTransport(ObjectMapper objectMapper,
                                                             McpServerProperties serverProperties) {
        return new WebMvcSseServerTransport(objectMapper, serverProperties.getSseMessageEndpoint());
    }

    @Bean
    public RouterFunction<ServerResponse> mvcMcpRouterFunction(WebMvcSseServerTransport transport) {
       return transport.getRouterFunction().filter(new HandlerFilterFunction<ServerResponse, ServerResponse>() {
            @Override
            public ServerResponse filter(ServerRequest request, HandlerFunction<ServerResponse> next) throws Exception {
                String path=request.path();
                System.out.println("获取鉴权");
                if (path.equals("sse")) {
                 //获取头部的token
                    Optional<String> token = request.param("token");
                    if (token.isPresent()){
                        //鉴权
                        if (token.get().equals("123456")){
                            return next.handle(request);
                        }else {
                            return ServerResponse.status(HttpStatus.UNAUTHORIZED).build();
                        }
                    }else {
                        return ServerResponse.status(HttpStatus.UNAUTHORIZED).build();
                    }
                }
                return next.handle(request);
            }
        });
    }
}
